﻿// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

#if !NET40
using System;
using System.IO;
using System.Buffers;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using CuteAnt.Buffers;
#if !NETCOREAPP_2_0_GREATER
using CuteAnt.AsyncEx;
#endif
#if DESKTOPCLR
using CuteAnt.Runtime;
#endif

namespace CuteAnt.IO.Pipelines
{
  internal static class StreamExtensions
  {
    public static Stream CreatePipelineStream(this IPipe pipe)
    {
      if (null == pipe) { throw new ArgumentNullException(nameof(pipe)); }

      return new PipelineStream(pipe);
    }

    public static Stream CreatePipelineWriterStream(this IPipeWriter writer, bool automaticFlushing = true)
    {
      if (null == writer) { throw new ArgumentNullException(nameof(writer)); }

      return new PipelineWriterStream(writer, automaticFlushing);
    }

    #region -- IPipeWriter --

    // We pick a value that is the largest multiple of 4096 that is still smaller than the large object heap threshold (85K).
    // The CopyTo/CopyToAsync buffer is short-lived and is likely to be collected at Gen0, and it offers a significant
    // improvement in Copy performance.
    private const int c_defaultCopyBufferSize = 81920;

    public static void CopyTo(this Stream stream, IPipeWriter writer, bool automaticFlushing = true)
    {
      if (null == writer) { throw new ArgumentNullException(nameof(writer)); }

#if DESKTOPCLR
      StreamToStreamCopy.Copy(stream, new PipelineWriterStream(writer, automaticFlushing), c_defaultCopyBufferSize, false);
#else
      stream.CopyTo(new PipelineWriterStream(writer, automaticFlushing), c_defaultCopyBufferSize);
#endif
    }

    /// <summary>Copies the content of a <see cref="Stream"/> into a <see cref="IPipeWriter"/>.</summary>
    public static Task CopyToAsync(this Stream stream, IPipeWriter writer, bool automaticFlushing = true, CancellationToken cancellationToken = default)
    {
      if (null == writer) { throw new ArgumentNullException(nameof(writer)); }

#if DESKTOPCLR
      return StreamToStreamCopy.CopyAsync(stream, new PipelineWriterStream(writer, automaticFlushing), c_defaultCopyBufferSize, false);
#else
      return stream.CopyToAsync(new PipelineWriterStream(writer, automaticFlushing), c_defaultCopyBufferSize, cancellationToken);
#endif
    }

    public static void CopyToEnd(this Stream stream, IPipeWriter writer, bool automaticFlushing = true)
    {
      try
      {
        stream.CopyTo(writer, automaticFlushing);
      }
      catch (Exception ex)
      {
        if (!automaticFlushing) { writer.FlushAsync().GetResult(); }
        writer.Complete(ex);
        return;
      }
      if (!automaticFlushing) { writer.FlushAsync().GetResult(); }
      writer.Complete();
    }

    public static async Task CopyToEndAsync(this Stream stream, IPipeWriter writer, bool automaticFlushing = true, CancellationToken cancellationToken = default)
    {
      try
      {
        await stream.CopyToAsync(writer, automaticFlushing, cancellationToken);
      }
      catch (Exception ex)
      {
        if (!automaticFlushing) { await writer.FlushAsyncAwaited(); }
        writer.Complete(ex);
        return;
      }
      if (!automaticFlushing) { await writer.FlushAsyncAwaited(); }
      writer.Complete();
    }

    #endregion

    #region -- ReadOnlyBuffer --

    /// <summary>Copies a <see cref="ReadOnlyBuffer<byte>"/> to a <see cref="Stream"/> asynchronously</summary>
    /// <param name="buffer">The <see cref="ReadOnlyBuffer<byte>"/> to copy</param>
    /// <param name="stream">The target <see cref="Stream"/></param>
    /// <returns></returns>
    public static void CopyTo(this ReadOnlyBuffer<byte> buffer, Stream stream)
    {
      if (buffer.IsSingleSegment)
      {
        WriteToStream(stream, buffer.First);
      }

      CopyMultipleToStream(buffer, stream);
    }

    /// <summary>Copies a <see cref="ReadOnlyBuffer<byte>"/> to a <see cref="Stream"/> asynchronously</summary>
    /// <param name="buffer">The <see cref="ReadOnlyBuffer<byte>"/> to copy</param>
    /// <param name="stream">The target <see cref="Stream"/></param>
    /// <returns></returns>
    public static Task CopyToAsync(this ReadOnlyBuffer<byte> buffer, Stream stream)
    {
      if (buffer.IsSingleSegment)
      {
        return WriteToStreamAsync(stream, buffer.First);
      }

      return CopyMultipleToStreamAsync(buffer, stream);
    }

    private static void CopyMultipleToStream(ReadOnlyBuffer<byte> buffer, Stream stream)
    {
      foreach (var memory in buffer)
      {
        WriteToStream(stream, memory);
      }
    }

    private static async Task CopyMultipleToStreamAsync(ReadOnlyBuffer<byte> buffer, Stream stream)
    {
      foreach (var memory in buffer)
      {
        await WriteToStreamAsync(stream, memory);
      }
    }

    private static void WriteToStream(Stream stream, ReadOnlyMemory<byte> readOnlyMemory)
    {
      if (MemoryMarshal.TryGetArray(readOnlyMemory, out ArraySegment<byte> data))
      {
        stream.Write(data.Array, data.Offset, data.Count);
      }
      else
      {
        // Copy required
        var pool = BufferManager.Shared;
        if (readOnlyMemory.Span.TryGetArray(pool, out data))
        {
          stream.Write(data.Array, data.Offset, data.Count);
          pool.Return(data.Array);
        }
      }
    }

    private static async Task WriteToStreamAsync(Stream stream, ReadOnlyMemory<byte> readOnlyMemory)
    {
      if (MemoryMarshal.TryGetArray(readOnlyMemory, out ArraySegment<byte> data))
      {
        await stream.WriteAsync(data.Array, data.Offset, data.Count)
            .ConfigureAwait(continueOnCapturedContext: false);
      }
      else
      {
        // Copy required
        var pool = BufferManager.Shared;
        if (readOnlyMemory.Span.TryGetArray(pool, out data))
        {
          await stream.WriteAsync(data.Array, data.Offset, data.Count).ConfigureAwait(continueOnCapturedContext: false);
          pool.Return(data.Array);
        }
      }
    }

    #endregion

    #region -- IPipeReader --

    //private const int c_defaultBufferSize = 1024 * 4;

    public static async Task CopyToAsync(this IPipeReader input, Stream stream, CancellationToken cancellationToken)
    {
      while (!cancellationToken.IsCancellationRequested)
      {
        var result = await input.ReadAsync();
        var inputBuffer = result.Buffer;
        try
        {
          if (inputBuffer.IsEmpty && result.IsCompleted)
          {
            return;
          }

          await inputBuffer.CopyToAsync(stream);
        }
        finally
        {
          input.Advance(inputBuffer.End);
        }
      }
    }

    public static async Task CopyToEndAsync(this IPipeReader input, Stream stream, CancellationToken cancellationToken)
    {
      try
      {
        await input.CopyToAsync(stream, cancellationToken);
      }
      catch (Exception ex)
      {
        input.Complete(ex);
        return;
      }
      return;
    }

    #endregion
  }
}
#endif
